"""
使用sglang调用DeepSeek Reasoner API实现:
1. 根据问题找到相关章节
2. 根据章节内容生成思维链和答案

使用方法:
python document_qa_with_reasoning.py
"""

import re
import json
import time
import sglang as sgl
from typing import Dict
from data_utils import load_doc_structure, match_chapter_ids_in_response, load_chapter_main

# 环境变量设置
QUESTIONS_FILE = "questions.json"  # 问题文件
DOC_STRUCTURE_FILE = "doc_structure.json"  # 文档结构文件
CHAPTERS_DIR = "./chapter_contents/"  # 章节文件目录（默认文件名为1.json, 2.json, 3.json, etc.）
OUTPUT_FILE = "think_ans.json"  # 输出文件
MAX_QUESTION_NUM = 30

MAX_RETRIES = 3  # 最大重试次数
RETRY_DELAY = 5  # 重试延迟(秒)

API = "deepseek-r1"

if API == "deepseek-r1":
    API_MODEL_NAME = "deepseek-reasoner"
    API_KEY = "sk-c80e688b2d544f16825ceca256d3c6e2"  # API密钥
    API_BASE_URL = "https://api.deepseek.com"  # API基础URL
    STREAM = False
elif API == "huoshan-deepseek-r1":
    API_MODEL_NAME = "deepseek-r1-250120"
    API_KEY = "9b1afd3f-bbe8-493f-9ba4-c17784556f0c"  # API密钥
    API_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"  # API基础URL
    STREAM = False
elif API == "huoshan-deepseek-r1-distill-qwen-7b":
    API_MODEL_NAME = "deepseek-r1-distill-qwen-7b-250120"
    API_KEY = "9b1afd3f-bbe8-493f-9ba4-c17784556f0c"  # API密钥
    API_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"  # API基础URL
    STREAM = False
elif API == "qwq-official":
    API_MODEL_NAME = "qwq-32b"
    API_KEY = "sk-8db036d3267f461bb042c9c9b14d8143"
    API_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
    STREAM = True


# 设置API后端
def setup_api():
    sgl.set_default_backend(
        sgl.OpenAI(
            API_MODEL_NAME,
            base_url=API_BASE_URL,
            api_key=API_KEY,
            is_chat_model=True
        )
    )


# 加载问题
def load_questions(questions_file: str) -> Dict[str, str]:
    with open(questions_file, 'r', encoding='utf-8') as f:
        return json.load(f)


# 查找与问题相关的章节，带异常重试
def find_related_sections_with_retry(question: str, sections_info: str) -> str:
    for attempt in range(MAX_RETRIES):
        try:
            @sgl.function
            def find_related_sections(s):
                s += sgl.system(
                    f"你是一个精确的文档检索助手。你的任务是从给定章节列表中找出与问题最相关的章节，严格避免引用不存在的章节。\n\n"
                    f"下面是文档的完整章节列表及其摘要：\n{sections_info}\n\n"
                    f"请注意：\n"
                    f"1. 只能推荐上述列表中实际存在的章节\n"
                    f"2. 不要编造或推测任何未列出的章节\n"
                    f"3. 如果问题与任何章节都不相关，请直接说明\"没有找到相关章节\"\n\n"
                    f"请按以下格式严格回答：\n"
                    f"\\chap{{实际章节号}}\\chap{{实际章节号}}\n"
                    f"简要解释这些章节与问题的相关性（2-3句话）\n\n"
                    f"示例回答：\n"
                    f"\\chap{{4.2}}\\chap{{5.3}}\n"
                    f"这些章节分别介绍了计算单元的基本原理和硬件架构抽象，直接回答了您关于处理器设计的问题。\n\n"
                )

                s += sgl.user(f"问题: {question}")

                s += sgl.assistant(sgl.gen(
                    name="result",
                    max_tokens=512,
                    temperature=0.2
                ))

            if STREAM:
                state = find_related_sections.run(stream=True)["result"]
            else:
                state = find_related_sections.run()["result"]

            return state
        except Exception as e:
            if attempt < MAX_RETRIES - 1:
                print(f"章节匹配失败，错误信息: {e}，将在 {RETRY_DELAY} 秒后重试...")
                time.sleep(RETRY_DELAY)
            else:
                print(f"章节匹配失败 {MAX_RETRIES} 次，使用默认章节。错误信息: {e}")
                return "\\chap{1} 由于匹配失败，默认使用Ascend C简介章节"


# 根据章节内容生成思维链和答案，带异常重试
def generate_cot_answer_with_retry(question: str, related_chapter_ids: list, chapter_contents: str) -> Dict[str, str]:
    for attempt in range(MAX_RETRIES):
        try:
            @sgl.function
            def generate_chain_of_thought_answer(s):
                s += sgl.system(
                    "请基于给定的问题和答案，生成详细的思维链过程。思维链应展示从问题到答案的完整推理路径，包括关键思考步骤、逻辑分析和推导过程。\n"
                    "思维链应当："
                    "1.展示解决问题的清晰思路;"
                    "2.包含必要的分析步骤和中间推理;"
                    "3.逐步引导至最终答案;"
                    "4.使用自然、流畅的语言;"
                    "5.避免使用\"根据原文\"、\"题目中提到\"等引用性表述。\n"

                    "示例格式：\n"
                    "问题：double buffer技术如何在Ascend C中应用？它如何提高计算效率？\n"
                    "思维链：[嗯，用户问的是Ascend C中怎么应用双缓冲技术以及它是如何提高计算效率的。首先，我得先回顾一下双缓冲技术的基本概念。双缓冲主要是通过交替使用两个缓冲区来隐藏数据传输和计算的延迟，这样可以让数据传输和计算并行进行。这在GPU或者NPU的编程中比较常见，用来提高数据吞吐量和计算效率。那Ascend C是华为的昇腾AI处理器的编程语言吧？所以用户可能是在使用昇腾芯片进行开发时遇到了性能优化的问题，想通过双缓冲来提升效率。我需要先确认Ascend C中的一些特性，比如是否支持异步的数据传输，或者是否有特定的内存管理指令，或者流水线并行的机制。比如，在CUDA中，可以使用cudaMemcpyAsync来实现异步传输，结合多流来隐藏延迟，昇腾的Ascend C是否有类似的机制呢？然后，用户想知道如何应用这个技术。应用的话，可能需要分步骤说明。比如，首先创建两个缓冲区，然后在一个缓冲区进行数据传输的同时，另一个缓冲区进行数据处理。交替使用这两个缓冲区，使得数据传输和计算重叠。不过在Ascend C中具体的实现方式可能涉及到任务并行、异步DMA传输命令或者特定的流水线指令。比如，可能需要在代码中启动异步的数据搬运任务，同时在另一个任务中进行计算，然后同步这两个任务。接下来，关于如何提高效率，双缓冲的主要优势在于减少了空闲时间。传统单缓冲的情况下，计算单元需要等待数据传输完成才能开始处理，而双缓冲可以让计算单元在处理当前数据块的同时，下一个数据块已经在传输中。这样就减少了整体的等待时间，提升了硬件利用率。比如，当计算单元在处理Buffer A的时候，DMA在传输Buffer B的数据，两者同时进行，这样每个块的总体时间近似于max(传输时间，计算时间)，而不是两者的总和。不过，在Ascend C的具体实现中，需要注意哪些方面呢？比如，内存对齐的要求，或者需要特定的API来管理缓冲区的切换。还可能涉及到任务之间的同步机制，比如事件或信号量，来确保数据已经传输完成才能开始计算。另外，可能还需要考虑数据的分块大小，以确保每个块的处理时间与传输时间尽可能接近，从而最大化效率。用户可能已经了解双缓冲的基本概念，但需要具体的代码示例或者步骤来说明如何在Ascend C中实现。比如，是否使用特定的数据搬运指令，或者如何分配和管理两个缓冲区。另外，可能还需要考虑昇腾处理器的架构特点，比如计算单元和存储之间的带宽，或者是否存在多级缓存，这些都可能影响双缓冲的实现方式和效果。可能存在的误区是认为双缓冲可以无限提高性能，但实际上如果数据传输时间和计算时间差距太大，可能无法完全隐藏延迟。比如，如果传输时间远大于计算时间，那么双缓冲可能只能部分优化，这时候可能需要更大的缓冲区或者更细粒度的任务划分。因此，在Ascend C中应用双缓冲时，需要根据具体的硬件特性和算法需求来调整缓冲区的大小和任务调度策略。总结一下，回答应该包括以下几点：双缓冲的基本原理，在Ascend C中的具体实现步骤（比如分配双缓冲区，异步数据搬运，计算与传输的重叠，同步机制），以及如何通过重叠计算和数据传输来提高效率。同时可能需要提到昇腾处理器的特定API或编程模型，如使用异步任务、流水线并行指令等。]\n"
                    "答案：[执行于AI Core上的指令队列主要包括如下几类，即Vector指令队列、Cube指令队列和MTE指令队列。不同指令队列间的相互独立性和可并行执行性，是double buffer优化机制的基石。矢量计算CopyIn、CopyOut过程使用MTE指令队列（MTE2、MTE3），Compute过程使用Vector指令队列（V），意味着CopyIn、CopyOut过程和Compute过程是可以并行的。如图13-1所示，考虑一个完整的数据搬运和计算过程，CopyIn过程将数据从GlobalMemory搬运到Local Memory，Vector计算单元完成计算后，经过CopyOut过程将计算结果搬回Global Memory。在此过程中，数据搬运与Vector计算串行执行，Vector计算单元无可避免存在资源闲置问题。举例而言，若CopyIn、Compute、CopyOut三阶段分别耗时t，则Vector的时间利用率仅为1/3，等待时间过长，Vector利用率严重不足。为减少Vector等待时间，double buffer机制将待处理的数据一分为二，比如Tensor1、Tensor2。如图13-2所示，当Vector对Tensor1中数据进行Compute时，Tensor2可以执行CopyIn的过程；而当Vector切换到计算Tensor2时，Tensor1可以执行CopyOut的过程。由此，数据的进出搬运和Vector计算实现并行执行，Vector闲置问题得以有效缓解。总体来说，double buffer是基于MTE指令队列与Vector指令队列的独立性和可并行性，通过将数据搬运与Vector计算并行执行以隐藏数据搬运时间并降低Vector指令的等待时间，最终提高Vector单元的利用效率，您可以通过为队列申请内存时设置内存块的个数来实现数据并行。]\n"

                    "请确保思维链逻辑严密，能够自然地连接问题和答案，让人能够理解从问题到答案的推导过程。"

                    "输出分成思维链和答案两部分内容，以如下格式输出：\n"
                    "<think>思维链<think>\n"
                    "<answer>答案<answer>\n"

                    "下面是一些文档章节内容，你可以参考它们来辅助你的思考，但记住，你的回答应该是基于你自己作为专家的思考而非文档的摘抄:"
                    f"\n\n{chapter_contents}\n\n"
                )

                s += sgl.user(f"遵循上面的指令，当前我的问题是: {question}")

                s += sgl.assistant(sgl.gen(
                    name="result",
                    max_tokens=2048,
                    temperature=0.7
                ))

            if STREAM:
                thinking_and_answer = generate_chain_of_thought_answer.run(stream=True)["result"]
            else:
                thinking_and_answer = generate_chain_of_thought_answer.run()["result"]

            # 分离思维链和答案
            # 通常格式为：思考过程部分...然后最终答案部分
            # thinking_and_answer = response.split("ASSISTANT:", 1)[1]
            if API == "huoshan-deepseek-r1":
                thinking = re.search(r'思维链：\s*(.*?)\s*答案：', thinking_and_answer, re.DOTALL).group(1).strip()
                answer = re.search(r'答案：\s*(.*)', thinking_and_answer, re.DOTALL).group(1).strip()
            else:
                thinking = re.search(r'<think>\s*(.*?)\s*</think>', thinking_and_answer, re.DOTALL).group(1).strip()
                answer = re.search(r'<answer>\s*(.*?)\s*</answer>', thinking_and_answer, re.DOTALL).group(1).strip()

            return {"question": question, "related_chapters": related_chapter_ids, "think": thinking, "answer": answer}
        except Exception as e:
            if attempt < MAX_RETRIES - 1:
                print(f"生成思维链和答案失败，错误信息: {e}，将在 {RETRY_DELAY} 秒后重试...")
                time.sleep(RETRY_DELAY)
            else:
                print(f"生成思维链和答案失败 {MAX_RETRIES} 次。错误信息: {e}")
                return {
                    "question": question,
                    "related_chapters": related_chapter_ids,
                    "think": f"由于API调用失败 {MAX_RETRIES} 次，无法生成思维链。错误信息: {str(e)}",
                    "answer": "由于技术原因，无法生成答案。"
                }

    # 这里不应该被执行到
    return {"think": "生成失败", "answer": "生成失败"}


# 主函数处理问题列表
def process_questions():
    # 设置API
    setup_api()

    # 加载问题
    print("Loading questions...")
    questions = load_questions(QUESTIONS_FILE)

    # 加载文档结构
    print("Loading doc structures...")
    sections_info = load_doc_structure(DOC_STRUCTURE_FILE)

    # 结果字典
    results = {}

    # 处理每个问题
    total_questions = len(questions)
    num_processed_questions = 0
    for question_id, question in questions.items():
        if num_processed_questions >= MAX_QUESTION_NUM:
            break

        print(f"处理问题 {question_id}/{total_questions}: {question[:50]}...")

        # 步骤1: 找到相关章节
        print(f"    try to find related sections...")
        related_sections_response = find_related_sections_with_retry(question, sections_info)
        related_chapter_ids = match_chapter_ids_in_response(related_sections_response, question_id)

        # 步骤2: 加载相关章节内容
        print(f"    try to load related sections...")
        all_chapter_contents = load_chapter_main(related_chapter_ids, CHAPTERS_DIR)

        # 步骤3: 生成思维链和答案
        print(f"    try to generate cot answer...")
        think_answer = generate_cot_answer_with_retry(question, related_chapter_ids, all_chapter_contents)

        # 保存结果到字典
        results[question_id] = think_answer

        num_processed_questions += 1

        # 定期保存结果以防中断
        if int(question_id) % 5 == 0 or int(question_id) == total_questions:
            # 保存阶段性结果
            with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            print(f"已完成 {question_id}/{total_questions} 个问题的处理，阶段性结果已保存")

    # 保存最终结果
    with open(OUTPUT_FILE, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)

    print(f"所有 {total_questions} 个问题处理完成，结果已保存到 {OUTPUT_FILE}")


def generate_questions(num_questions=100):
    setup_api()
    sections_info = load_doc_structure(DOC_STRUCTURE_FILE)

    @sgl.function
    def generate_questions(s):
        s += sgl.system(
            "你是一位精通文本分析的AI助手。我将提供一段文档章节，你的任务是深入理解这段内容，并从中提取出100个不重复的问题。"
            "请遵循以下要求："
            "1. 仔细阅读和理解所提供的文档内容，把握其核心观点、关键概念和逻辑结构。"
            "2. 基于文档内容生成100个有价值的问题，这些问题应该："
            "  - 覆盖文档的主要内容和细节"
            "  - 包含不同认知层次的问题（如理解、分析、应用、评价等）"
            "  - 避免表述模糊或过于宽泛的问题"
            "  - 确保每个问题都基于文档中的实际内容"
            "3. 所有问题必须使用以下格式严格输出："
            "  <q>问题1</q>"
            "  <q>问题2</q>"
            "  <q>问题3</q>"
            "  ..."
            "  <q>问题100</q>"
            "4. 确保100个问题之间不存在重复或高度相似的问题，每个问题应该针对文档的不同方面或角度。"
            "5. 如果文档内容有限，也请尽力从不同角度提出问题，包括推理性、假设性问题，但这些问题仍应与文档内容相关。"
            "请直接生成符合要求的100个问题，不需要其他解释或说明。"
        )

        s += sgl.user(f"遵循上面的指令，当前我的文档是: {sections_info}")

        s += sgl.assistant(sgl.gen(
            name="result",
            max_tokens=2048,
            temperature=0.7
        ))

    if STREAM:
        generated_questions = generate_questions.run(stream=True)["result"]
    else:
        generated_questions = generate_questions.run()["result"]

    print(f"generated_questions = {generated_questions}")


if __name__ == "__main__":
    # generate_questions()
    process_questions()